/*
主天线在前，从天线在后
*/

# pragma once

# include <homo_ocp.h>

# include "cpp_tools.h"
# include "udpSettings.h"
# include "estimator_base.h"

# ifdef INPUT_REAL
# include "filter_W1M1_C.h"
# include "WHEELTEC_N100.h"
# include "WTRTK_982_4G.h"
# endif

class estimator_RTK_IMU : public estimator_base
{
private:
    struct mtData
    {
        cpp_tools::mtDataManager<std::list<Eigen::Vector4d>> IMU_list; // t, a_x, a_y, a_z
        cpp_tools::mtDataManager<std::list<Eigen::Vector4d>> AHRS_list; // t, psi, gamma, omega
        cpp_tools::mtDataManager<Eigen::Vector<double, 6>> RTK_data; // t, x_m, y_m, x_s, y_s, varphi

        cpp_tools::mtDataManager<Eigen::Vector<double, 10>> state; // 当前最优估计 x, y, varphi, psi, gamma, v_x, v_y, omega, a_x, a_y
        cpp_tools::mtDataManager<double> t_update = -1;

        cpp_tools::mtDataManager<std::tuple<double,double>> position_cor = {0, 0}; // x_cor, y_cor
    } m_mtData;

    // =============================定义参数=============================
    struct frequency // 频率，Hz
    {
        # ifdef INPUT_REAL
        double imu_port = 1e4; // 端口读取频率
        double rtk_port = 1e4; // 端口读取频率
        # else
        double state = 1e4;
        # endif
    } const m_frequency;

    const double gravity = 9.82; // 重力加速度
    const double pi = cpp_tools::pi; // 圆周率
    const double m_INF = std::numeric_limits<double>::infinity();

    // RTK
    std::string RTK_device = "/dev/ttyACM3";
    int RTK_bauds = 115200;
    bool RTK_ready = false; // 是否已经收到RTK信号
    
    const double dx_front = 0.216; // 前面的天线到质心的距离
    const double dx_rear = -0.258; // 后面的天线到质心的距离

    double RTK_lastUpdateTime = cpp_tools::getCurrentTimeAsDouble();
    Eigen::Vector2d RTK_lastPose;
    Eigen::Vector<double, 5> RTK_opt;

    // IMU
    std::string IMU_device = "/dev/ttyUSB0";
    int IMU_bauds = 921600;
    const double maxDis = 5; // 加速度错误值滤波的最大限差

    // ==================================================================
    cpp_tools::threadPool m_threadPool; // 线程池

    homo_ocp::prob m_prob;
    std::unique_ptr<homo_ocp::solver> m_solver;

    size_t listMaxSize = 100;

    bool isRunning = true;

    cpp_tools::clock debug_clock;

    // ==================================================================

    /// @brief 消除重力加速度作用
    /// @param imu 
    /// @param ahrs 
    void rmG(Eigen::Vector3d & imu, const Eigen::Vector3d & ahrs)
    {
        double psi = ahrs(0);
        double gamma = ahrs(1);
        double omega = ahrs(2);

        Eigen::VectorXd g(3);
        g << 0, 0, gravity;
        Eigen::VectorXd da 
            = cpp_tools::Euler_Rotation_X(gamma).inverse()*
            cpp_tools::Euler_Rotation_Y(psi).inverse()*g;
        imu -= da;
    }

    void IMU_forward()
    {
        // std::cout << "update using IMU data" << std::endl;

        auto ac_state = m_mtData.state.write();

        double varphi = (*ac_state.access())(2);
        double psi = (*ac_state.access())(3);
        double gamma = (*ac_state.access())(4);
        double v_x = (*ac_state.access())(5);
        double v_y = (*ac_state.access())(6);

        auto imu_list = *m_mtData.IMU_list.read().access();

        // 计算对地面的投影速度
        Eigen::VectorXd v(3);
        v << v_x,v_y,0;
        Eigen::MatrixXd R = cpp_tools::Euler_Rotation_Z(varphi)*cpp_tools::Euler_Rotation_Y(psi)*cpp_tools::Euler_Rotation_X(gamma);
        Eigen::VectorXd v_prj = R*v;

        // 积分
        if (imu_list.size() < 2)
        {
            return;
        }
        else
        {
            auto imu_1 = imu_list.back();
            imu_list.pop_back();
            auto imu_0 = imu_list.back();

            double Delta_t = imu_1(0) - imu_0(0);
            double a_x = imu_1(1);
            double a_y = imu_1(2);
            double omega = (*ac_state.access())(7);

            (*ac_state.access())(0) += v_prj(0)*Delta_t;
            (*ac_state.access())(1) += v_prj(1)*Delta_t;
            (*ac_state.access())(5) += (a_x + omega*v_y)*Delta_t;
            (*ac_state.access())(6) += (a_y - omega*v_x)*Delta_t;
            (*ac_state.access())(8) = a_x;
            (*ac_state.access())(9) = a_y;
        }
    }

    void AHRS_forward()
    {
        // std::cout << "update using AHRS data" << std::endl;

        auto ac_state = m_mtData.state.write();
        
        auto ahrs_list = *m_mtData.AHRS_list.read().access();
        if (ahrs_list.size() < 2)
        {
            return;
        }
        else
        {
            auto ahrs_1 = ahrs_list.back();
            ahrs_list.pop_back();
            auto ahrs_0 = ahrs_list.back();

            double Delta_t = ahrs_1(0) - ahrs_0(0);
            double omega = ahrs_1(3);

            (*ac_state.access())(2) += omega*Delta_t;
            (*ac_state.access())(7) = omega;
        }
    }

    void RTK_update()
    {
        // std::cout << "update using RTK data" << std::endl;

        auto ac_state = m_mtData.state.write();

        auto rtk_pkg = *m_mtData.RTK_data.read().access();
        Eigen::Vector3d rtk;
        rtk(0) = (rtk_pkg(1)*dx_rear - rtk_pkg(3)*dx_front)/(dx_rear - dx_front);
        rtk(1) = (rtk_pkg(2)*dx_rear - rtk_pkg(4)*dx_front)/(dx_rear - dx_front);
        rtk(2) = rtk_pkg(5);

        Eigen::Vector3d imu_avg;
        Eigen::Vector3d ahrs_avg;
        double Delta_t = 0;
        {
            auto ac_imu_list = m_mtData.IMU_list.write();
            auto ac_ahrs_list = m_mtData.AHRS_list.write();

            double lastUpdateTime = RTK_lastUpdateTime;
            double updateTime = rtk_pkg(0);

            RTK_lastUpdateTime = updateTime;

            // 初始化
            if (!RTK_ready)
            {
                RTK_lastPose(0) = rtk(0);
                RTK_lastPose(1) = rtk(1);

                RTK_opt(0) = rtk(0);
                RTK_opt(1) = rtk(1);
                RTK_opt(2) = rtk(2);

                RTK_ready = true;

                return;
            }

            Delta_t = updateTime - lastUpdateTime; // 单位 s

            std::function<void (std::list<Eigen::Vector4d> &, Eigen::Vector3d &)> 
            listIntV4d([lastUpdateTime, updateTime](std::list<Eigen::Vector4d> & list, Eigen::Vector3d & vInt) -> void {
                double t_last = lastUpdateTime;
                while (!list.empty())
                {
                    auto v4d = list.front();
                    double t = v4d(0);
                    if (t < t_last)
                    {
                        list.pop_front();
                    }
                    else if (t <= updateTime)
                    {
                        vInt += v4d.bottomRows(3)*(t - t_last);

                        t_last = t;
                        
                        list.pop_front();

                        if (list.empty())
                        {
                            vInt += v4d.bottomRows(3)*(updateTime - t_last);

                            v4d(0) = updateTime;
                            list.push_back(v4d);

                            break;
                        }
                    }
                    else
                    {
                        vInt += v4d.bottomRows(3)*(updateTime - t_last);

                        break;
                    }
                }
            });

            Eigen::Vector3d imu_int;
            imu_int.setZero();
            listIntV4d(*ac_imu_list.access(), imu_int);
            imu_avg = imu_int/Delta_t;
        
            Eigen::Vector3d ahrs_int;
            ahrs_int.setZero();
            listIntV4d(*ac_ahrs_list.access(), ahrs_int);
            ahrs_avg = ahrs_int/Delta_t;
        }

        if (0)
        {
            double psi = ahrs_avg(0);
            double gamma = ahrs_avg(1);
            double omega = ahrs_avg(2);

            double a_x = imu_avg(0);
            double a_y = imu_avg(1);

            double x_bias = rtk(0);
            double y_bias = rtk(1);
            double varphi = rtk(2);

            rtk(0) -= x_bias;
            rtk(1) -= y_bias;

            RTK_opt(0) -= x_bias;
            RTK_opt(1) -= y_bias;

            // 因为引入了航向角的运动学模型，因此需要保证mod2pi不会导致数值突变
            while (std::abs(varphi - RTK_opt(2)) > pi)
            {
                if (varphi - RTK_opt(2) > pi)
                {
                    RTK_opt(2) += 2*pi;
                }
                else
                {
                    RTK_opt(2) -= 2*pi;
                }
            }
            
            // 权重
            Eigen::Vector<double, 3> Q;
            Eigen::Vector<double, 5> R;
            Eigen::Vector<double, 1> D;
            Q << 4,4,4;
            R << 1,1,1,0.1,0.1;
            D << 0;

            double diff_omega = std::abs(omega) + 0.4;
            
            // 问题的初值
            Eigen::VectorXd r_0(m_prob.n_x);
            r_0 << RTK_opt, Eigen::VectorXd::Zero(5);

            // 计算数据的间隔时间，填充问题并求解
            m_prob.p << 
                Delta_t, // 参数中的时间单位是 s
                RTK_opt, // 上一时刻更新的最优估计
                rtk, // 当前rtk更新
                ahrs_avg, // 参考姿态
                imu_avg.topRows(2), // 时间段内的等效加速度
                Q,
                R,
                D;

            m_prob.lbx << -0.2, -0.2, -m_INF, -m_INF, -m_INF, -m_INF, -m_INF, -Delta_t*diff_omega, -m_INF, -m_INF;
            m_prob.ubx << 0.2, 0.2, m_INF, m_INF, m_INF, m_INF, m_INF, Delta_t*diff_omega, m_INF, m_INF;

            m_solver.get()->reLoadBoundary(&m_prob.lbx,&m_prob.ubx,1);
            m_solver.get()->reLoadParam(m_prob.p);

            auto sol = m_solver.get()->solve(r_0);

            RTK_opt = sol.x.topRows(5);
            RTK_opt(0) += x_bias;
            RTK_opt(1) += y_bias;
            while (std::abs(RTK_opt(2)) > pi)
            {
                if (RTK_opt(2) > pi)
                {
                    RTK_opt(2) -= 2*pi;
                }
                else
                {
                    RTK_opt(2) += 2*pi;
                }
            }
 
            auto [x_cor, y_cor] = *m_mtData.position_cor.read().access();
            (*ac_state.access())(0) = RTK_opt(0) - x_cor;
            (*ac_state.access())(1) = RTK_opt(1) - y_cor;
            (*ac_state.access())(2) = RTK_opt(2);
            (*ac_state.access())(5) = RTK_opt(3);
            (*ac_state.access())(6) = RTK_opt(4);
        }
        else // 更新方式二
        {
            double v_x = (rtk(0) - RTK_lastPose(0))/Delta_t;
            double v_y = (rtk(1) - RTK_lastPose(1))/Delta_t;
            Eigen::MatrixXd R = cpp_tools::Euler_Rotation_Z(rtk(2))*cpp_tools::Euler_Rotation_Y(ahrs_avg(0))*cpp_tools::Euler_Rotation_X(ahrs_avg(1));
            Eigen::Vector3d v_prj(v_x,v_y,0);
            Eigen::Vector3d v = R.transpose()*v_prj;
            
            auto [x_cor, y_cor] = *m_mtData.position_cor.read().access();
            (*ac_state.access())(0) = rtk(0) - x_cor;
            (*ac_state.access())(1) = rtk(1) - y_cor;
            (*ac_state.access())(2) = rtk(2);
            (*ac_state.access())(5) = v(0);
            (*ac_state.access())(6) = v(1);
    
            RTK_lastPose(0) = rtk(0);
            RTK_lastPose(1) = rtk(1);
        }
    }

public:
    estimator_RTK_IMU()
        : estimator_base()
    {
        std::cout << "state estimator start ..." << std::endl;

        // 问题初始化
        HOMO_OCP_PROB(m_prob,filter_W1M1);
        if (m_prob.n_x != nz_h[0] || m_prob.n_x != nz_A[1] || m_prob.n_eq != nz_A[0])
        {
            throw std::runtime_error("问题维度不一致，请检查提供问题及其梯度信息的文件");
        }

        m_solver = std::make_unique<homo_ocp::solver>(m_prob,m_INF,2,0);

        # ifdef INPUT_REAL
        cpp_tools::XmlParameterManager m_xml("device", host_path + settings_path + "/device/settings.xml");
        bool writeXML = false;
        if (!m_xml.getParameter("IMU", "fileName", IMU_device))
        {
            std::cout << "IMU设备路径使用使用默认值: " << IMU_device << std::endl;
            m_xml.setParameter("IMU", "fileName", IMU_device);

            writeXML = true;
        }
        if (!m_xml.getParameter("IMU", "bauds", IMU_bauds))
        {
            std::cout << "IMU波特率使用使用默认值: " << IMU_bauds << std::endl;
            m_xml.setParameter("IMU", "bauds", IMU_bauds);

            writeXML = true;
        }
        if (!m_xml.getParameter("RTK", "fileName", RTK_device))
        {
            std::cout << "RTK设备路径使用使用默认值: " << RTK_device << std::endl;
            m_xml.setParameter("RTK", "fileName", RTK_device);

            writeXML = true;
        }
        if (!m_xml.getParameter("RTK", "bauds", RTK_bauds))
        {
            std::cout << "RTK波特率使用使用默认值: " << RTK_bauds << std::endl;
            m_xml.setParameter("RTK", "bauds", RTK_bauds);

            writeXML = true;
        }
        if (writeXML)
        {
            m_xml.saveToFile(host_path + settings_path + "/device/settings.xml");
        }
        
        // =============================imu 端口读取线程=============================
        std::thread t_imu_port([&]() -> void {
            WHEELTEC_N100::driver m_imu(IMU_device.data(),IMU_bauds);

            Eigen::VectorXd imu_last(3); // 记录上一个数据
            while (isRunning)
            {
                auto r = m_imu.reciveData();
                double t_read = cpp_tools::getCurrentTimeAsDouble();
                if (r == 1)
                {
                    auto imu_pkg = m_imu.get_IMUdata();

                    // 滤波
                    double a_x = imu_pkg->accelerometer_x;
                    double a_y = -imu_pkg->accelerometer_y;
                    double a_z = -imu_pkg->accelerometer_z;
                    Eigen::Vector3d imu_new;
                    imu_new << a_x,a_y,a_z;
                    if ((imu_new.array().abs() > 16).any() || (imu_new.array().isNaN()).any() || imu_new.array().isZero())
                    {
                        // std::cout << "IMU 疑似有错误读数" << imu_new.transpose() << std::endl;
                    }
                    else
                    {
                        {
                            auto ac_IMU = m_mtData.IMU_list.write();
    
                            size_t overDis = 0;
                            for (auto &&imu_old : *ac_IMU.access())
                            {
                                auto dis = imu_old.bottomRows(3) - imu_new;
                                if ((dis.array().abs() > maxDis).any())
                                {
                                    overDis += 1;
                                }
                            }
    
                            // 检查上一个值与此次的值是否一致，如果一致则不认为是异常值
                            bool acceptThis = false;
                            if (((imu_last - imu_new).array().abs() <= maxDis).all())
                            {
                                acceptThis = true;
                            }
    
                            // 更新
                            if (overDis < ac_IMU.access()->size() || acceptThis || ac_IMU.access()->size() == 0)
                            {
                                auto ahrs = m_mtData.AHRS_list.read().access()->back();
                                rmG(imu_new, ahrs.bottomRows(3));

                                Eigen::VectorXd tmpV(4);
                                tmpV << t_read, imu_new;
                                ac_IMU.access()->push_back(tmpV);

                                while (ac_IMU.access()->size() > listMaxSize)
                                {
                                    ac_IMU.access()->pop_front();
                                }
                            }
                        }

                        auto ac_t_update = m_mtData.t_update.write();
                        IMU_forward();
                        *ac_t_update.access() = t_read;

                        imu_last = imu_new;
                    }
                }
                else if (r == 2)
                {
                    auto ahrs_pkg = m_imu.get_AHRSdata();

                    double psi = -ahrs_pkg->Pitch;
                    double gamma = ahrs_pkg->Roll;
                    double omega = -ahrs_pkg->HeadingSpeed;
                    Eigen::VectorXd ahrs_new(3);
                    ahrs_new << psi,gamma,omega;
                    if ((ahrs_new.array().abs() > 16).any() || (ahrs_new.array().isNaN()).any() || ahrs_new.array().isZero())
                    {
                        // std::cout << "AHRS 疑似有错误读数: " << ahrs_new.transpose() << std::endl;
                    }
                    else
                    {
                        {
                            auto ac_AHRS = m_mtData.AHRS_list.write();

                            Eigen::VectorXd tmpV(4);
                            tmpV << t_read, ahrs_new;
                            ac_AHRS.access()->push_back(tmpV);

                            while (ac_AHRS.access()->size() > listMaxSize)
                            {
                                ac_AHRS.access()->pop_front();
                            }
                        }

                        auto ac_t_update = m_mtData.t_update.write();
                        AHRS_forward();
                        *ac_t_update.access() = t_read;
                    }
                }

                std::this_thread::sleep_for(std::chrono::microseconds(std::max((int)(1e6/m_frequency.imu_port),1)));
            }
        });
        m_threadPool.push_back(t_imu_port);

        // =============================rtk端口读取线程=============================
        std::thread t_rtk_port([&]() -> void {
            WTRTK_982_4G::driver m_rtk(RTK_device.data(),RTK_bauds);
            m_rtk.printINFO = 0;

            Eigen::VectorXd buffer = Eigen::VectorXd::Constant(5, std::numeric_limits<double>::quiet_NaN());
            while (isRunning)
            {
                auto messages = m_rtk.reciveData();

                double t_read = cpp_tools::getCurrentTimeAsDouble();
                for (auto &&msg : messages)
                {
                    if (msg == "GGA")
                    {
                        // printf("\nGGA lon: %.10f, lat: %.10f \n",m_driver.getGGA().lon,m_driver.getGGA().lat);

                        auto gga = m_rtk.getGGA();
                        double x;
                        double y;
                        cpp_tools::latLonToGaussKruger(gga.lat, gga.lon, x, y);
                        buffer.topRows(2) = Eigen::Vector2d(x,-y);

                        // printf("\nGGA x: %.10f, y: %.10f \n", x, y);
                    }
                    else if (msg == "GGAH")
                    {
                        // printf("\nGGAH lon: %.10f, lat: %.10f \n",m_driver.getGGA().lon,m_driver.getGGA().lat);

                        auto ggah = m_rtk.getGGAH();
                        double x;
                        double y;
                        cpp_tools::latLonToGaussKruger(ggah.lat, ggah.lon, x, y);
                        buffer.middleRows(2,2) = Eigen::Vector2d(x,-y);

                        // printf("\nGGAH x: %.10f, y: %.10f \n", x, y);
                    }
                    else if (msg == "HPR")
                    {
                        // printf("\nheading: %.4f (deg) \n",m_driver.getHPR().heading);

                        auto hpr = m_rtk.getHPR();
                        double varphi = -(hpr.heading - 180)/180*cpp_tools::pi;
                        buffer.bottomRows(1) = Eigen::Vector<double, 1>(varphi);

                        // printf("\nheading: %.10f (rad) \n", varphi);
                    }
                }

                // 只有全部都有数据时才更新
                if (!buffer.array().isNaN().any()) 
                {
                    {
                        auto ac_RTK = m_mtData.RTK_data.write();
    
                        Eigen::VectorXd tmpV(6);
                        tmpV << t_read, buffer;
                        *ac_RTK.access() = tmpV;

                        buffer = Eigen::VectorXd::Constant(5, std::numeric_limits<double>::quiet_NaN());
                    }

                    auto ac_t_update = m_mtData.t_update.write();
                    RTK_update();
                    *ac_t_update.access() = t_read;
                }

                std::this_thread::sleep_for(std::chrono::microseconds(std::min((int)(1e6/m_frequency.rtk_port),10)));
            }
        });
        m_threadPool.push_back(t_rtk_port);

        # else // INPUT_REAL
        // =============================从仿真环境中一次性获取数据=============================
        std::thread t_simState([&]() -> void {
            cpp_tools::udpVecReceiver<double> stateReceiver(udpSettings::PORT::lower_sim::state_len,udpSettings::PORT::lower_sim::state,1);

            cpp_tools::clock imu_clock;
            cpp_tools::clock rtk_clock;
            while (isRunning)
            {
                auto t = std::chrono::steady_clock::now();

                int info = stateReceiver.receive();
                if (info == 0)
                {
                    auto data = stateReceiver.get();
                    
                    double t_read = cpp_tools::getCurrentTimeAsDouble();
                    if (rtk_clock.toc(0) >= 50) // RTK 50ms 更新一次
                    {
                        // std::cout << "RTK update" << std::endl;

                        rtk_clock.tic();
                        
                        double x_front = data(0);
                        double y_front = data(1);
                        double x_rear = data(2);
                        double y_rear = data(3);
                        double varphi = atan2(y_front - y_rear, x_front - x_rear);
                        Eigen::Vector<double, 5> rtk_data(x_front, y_front, x_rear, y_rear, varphi);

                        Eigen::VectorXd tmpV(6);
                        tmpV << t_read, rtk_data;
                        *m_mtData.RTK_data.write().access() = tmpV;

                        auto ac_t_update = m_mtData.t_update.write();
                        RTK_update();
                        *ac_t_update.access() = t_read;
                    }
                    if (imu_clock.toc(0) >= 5) // IMU 5ms 更新一次
                    {
                        // std::cout << "IMU update" << std::endl;

                        imu_clock.tic();

                        Eigen::Vector3d imu = data.middleRows(4,3);
                        Eigen::Vector3d ahrs = data.bottomRows(3);

                        imu(1) *= -1;
                        imu(2) *= -1;

                        ahrs(0) *= -1;
                        ahrs(2) *= -1;

                        rmG(imu, ahrs);

                        Eigen::VectorXd tmpV(4);
                        tmpV << t_read, imu;
                        m_mtData.IMU_list.write().access()->push_back(tmpV);

                        tmpV << t_read, ahrs;
                        m_mtData.AHRS_list.write().access()->push_back(tmpV);

                        auto ac_t_update = m_mtData.t_update.write();
                        IMU_forward();
                        AHRS_forward();
                        *ac_t_update.access() = t_read;
                    }
                }
                
                std::this_thread::sleep_until(t + std::chrono::microseconds(std::max((int)(1e6/m_frequency.state),1)));
            }
        });
        m_threadPool.push_back(t_simState);

        # endif // INPUT_REAL
    };

    ~estimator_RTK_IMU()
    {
        shutdown();

        std::cout << "state estimator exit." << std::endl;
    };

    void shutdown()  override
    {
        isRunning = false;

        m_threadPool.joinAll();
    };

    Eigen::VectorXd getStateAsVector()  override
    {
        return *m_mtData.state.read().access();
    }

    state getStateAsStruct() override
    {
        auto stateVec = getStateAsVector();

        int i = 0;
        return {
            stateVec(i++),
            stateVec(i++),
            stateVec(i++),
            stateVec(i++),
            stateVec(i++),
            stateVec(i++),
            stateVec(i++),
            stateVec(i++),
            stateVec(i++),
            stateVec(i++)
        };
    };

    /// @brief 获取最近一次状态更新的时间，-1表示没有任何更新
    /// @return 
    double getUpdateTime() override
    {
        return *m_mtData.t_update.read().access();
    }

    /// @brief 增加rtk定位x, y, varphi的修正量
    /// @param x 
    /// @param y 
    void reInitFrame(double x, double y) override
    {
        auto [x_cor, y_cor] = *m_mtData.position_cor.read().access();
        x_cor += x;
        y_cor += y;
        *m_mtData.position_cor.write().access() = {x_cor, y_cor};
    }
};